package com.jacklin.snowy.limitstrategy.service.impl;

import com.jacklin.snowy.limitstrategy.constants.RedisKeyConstant;
import com.jacklin.snowy.limitstrategy.dto.RequestLimitDTO;
import com.jacklin.snowy.limitstrategy.enums.RequestLimitType;
import com.jacklin.snowy.limitstrategy.service.RequestLimitService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.List;
import java.util.UUID;

/**
 * 漏桶算法 限流
 *
 * @author: jacklin
 * @since: 2022/5/10 14:11
 */
@Slf4j
@Service
public class LeakyBucketRateLimitServiceImpl implements RequestLimitService {

    @Resource(name = "leakyBucketPopThreadPoolScheduler")
    private ThreadPoolTaskScheduler scheduler;

    @Value("${request-limit.scan-package}")
    private String scanPackage;

    @Autowired
    private ResourcePatternResolver resourcePatternResolver;

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Override
    public boolean checkRequestLimit(RequestLimitDTO requestLimitDTO) {
        // request:limit:qps:leakyBucket:tokenBucketTest
        String key = RedisKeyConstant.RequestLimit.QPS_LEAKY_BUCKET + requestLimitDTO.getKey();
        Long size = redisTemplate.opsForList().size(key);
        if (size != null && size >= requestLimitDTO.getLimiter().limitCount()) {
            log.info("【{}】限流控制，漏桶中容量已满，请被拦截", requestLimitDTO.getKey());
            return true;
        } else {
            log.info("【{}】漏桶中容量未满，当前水滴容量：{}，漏桶容量：{}", requestLimitDTO.getKey(), size, requestLimitDTO.getLimiter().limitCount());
            redisTemplate.opsForList().leftPush(key, UUID.randomUUID().toString());
            return false;
        }
    }

    /**
     * 定数流出令牌
     *
     * @author: jacklin
     * @date: 2022/5/10
     **/
    @PostConstruct
    public void popToken() {
        List<RequestLimitDTO> list = this.getTokenLimitList(resourcePatternResolver, RequestLimitType.LEAKY_BUCKET, scanPackage);
        if (list.isEmpty()) {
            log.info("未扫描到使用【漏桶限流】注解的方法，结束生成令牌线程");
            return;
        }
        list.forEach(requestLimitDTO -> {
            scheduler.scheduleAtFixedRate(() -> {
                /*
                 * 安排给定的Runnable ，尽快开始并在给定的时间段内调用它
                 * task – 触发器触发时执行的 Runnable
                 * period – 任务连续执行之间的间隔（以毫秒为单位）
                 **/
                String key = RedisKeyConstant.RequestLimit.QPS_LEAKY_BUCKET + requestLimitDTO.getKey();
                //截取List在start和end之间的元素处key列表
                redisTemplate.opsForList().trim(key, requestLimitDTO.getLimiter().limitPeriodCount(), -1);
                log.info("【{}】漏出{}个水滴", key, requestLimitDTO.getLimiter().limitPeriodCount());
            }, requestLimitDTO.getLimiter().period());
        });
    }

    @Override
    public RequestLimitType getType() {
        return RequestLimitType.LEAKY_BUCKET;
    }
}
